#!/usr/bin/python
# -*- coding: UTF-8 -*-
 
from xml.dom.minidom import parse
import xml.dom.minidom
import networkx as nx
import matplotlib.pyplot as plt

# 使用minidom解析器打开 XML 文档
DOMTree = xml.dom.minidom.parse("platform.xml")
collection = DOMTree.documentElement
if collection.hasAttribute("shelf"):
   print("Root element : %s" % collection.getAttribute("shelf"))
 
hosts = collection.getElementsByTagName("host")
hosts_name = []
for host in hosts:
    hosts_name.append(host.getAttribute('id'))
print("hosts: ", hosts_name)
links = collection.getElementsByTagName("link")
link_id = []
for link in links:
    link_id.append(link.getAttribute('id'))
print("links: ", link_id)
topo = []
#element: (src, dst), [4, 3, 2, 1] (links)
# print(link_id.index('3'))
print("==========Topo=============")
G = nx.Graph()
route = collection.getElementsByTagName('route')
dict = {}
tmp_node = 0
for host in hosts_name:
    G.add_node(host, weight=10)

for r in route:
    src = r.getAttribute('src')
    dst = r.getAttribute('dst')
    link_ctn = r.getElementsByTagName('link_ctn')
    links = []
    # G.add_edge(src, dst)
    for l in link_ctn:
        links.append(l.getAttribute('id'))
    print(src, dst, links)
    if len(links) == 1:
        G.add_edge(src, dst)
    else:
        for i in range(len(links)):
            if i != len(links)-1:
                dict.setdefault(links[i], {})
                dict[links[i]].setdefault(links[i+1], -1)
                dict.setdefault(links[i+1], {})
                dict[links[i+1]].setdefault(links[i], -1)
                if dict[links[i]][links[i+1]] == -1 and dict[links[i+1]][links[i]] == -1:
                    #no this node
                    G.add_node(str(tmp_node), weight=1)
                    dict[links[i]][links[i+1]] = tmp_node
                    dict[links[i+1]][links[i]] = tmp_node
                    tmp_node = tmp_node + 1
                if i == 0:
                    sc = src
                else:
                    sc = str(dict[links[i-1]][links[i]])

                if i == len(links)-1:
                    ds = dst
                else:
                    ds = str(dict[links[i]][links[i+1]])
            else:
                sc = str(dict[links[i-1]][links[i]])
                ds = dst
            G.add_edge(sc, ds)


nlarge=[n for (n,d) in G.nodes(data=True) if d['weight'] >1]
nsmall=[n for (n, d) in G.nodes(data=True) if d['weight'] <=1]
print(nlarge)
print(nsmall)
pos=nx.spring_layout(G)

label = {}
for node in hosts_name:
    label.setdefault(node, node)

nx.draw_networkx_nodes(G,pos, node_size=500, nodelist=nlarge, label=label)
nx.draw_networkx_nodes(G,pos, node_size=40, nodelist=nsmall)
nx.draw_networkx_labels(G, pos, labels=label, font_size=10)
nx.draw_networkx_edges(G, pos)
# nx.draw_networkx(G, pos=nx.spring_layout(G))
# nx.draw_networkx(G)
plt.savefig('./topo.png')
print(dict)